import uuid
from typing import List, Optional

from app.core.memory.memory_manager import get_session_history
from fastapi import APIRouter, Body
from pydantic import BaseModel, Field

from app.core.chains import get_conversation_chain

router = APIRouter()


class AgentChatRequest(BaseModel):
    input: str
    conversation_id: Optional[str] = Field(None,
                                           description="A unique ID for the conversation. If not provided, a new one will be generated.")
    tool_config: List[str] = Field(default=[], description="List of tool names to use, e.g., ['get_current_weather'].")


class AgentChatResponse(BaseModel):
    output: str
    conversation_id: str


@router.post("/", response_model=AgentChatResponse)
def agent_based_chat(request: AgentChatRequest = Body(...)):
    """
    Handles a chat request. If tools are specified, it uses an Agent.
    Otherwise, it falls back to a simple conversational LLM.
    """
    conv_id = request.conversation_id or str(uuid.uuid4())
    print(f"--- Agent/LLM Chat endpoint called for conversation_id: {conv_id} ---")

    chat_history_backend = get_session_history(conv_id)

    conversation_chain = get_conversation_chain(
        tool_names=request.tool_config,
        chat_history_backend=chat_history_backend
    )

    result = conversation_chain.invoke({"input": request.input})

    # Handle different output keys from AgentExecutor ('output') and LLMChain ('text')
    final_output = result.get("output", result.get("text"))

    return AgentChatResponse(
        output=final_output,
        conversation_id=conv_id
    )



